from langchain_openai import OpenAIEmbeddings
from langchain_chroma import Chroma
from langchain_openai import OpenAI
from langchain.chains import RetrievalQA

import os

os.environ['OPENAI_API_KEY'] = 'sk-ioXCJ3W2IQlpDzHu268aEeF7CcAf4b1f9f8c46A248C74a9d'
os.environ['OPENAI_BASE_URL'] = 'https://api.xiaoai.plus/v1'

embeddings = OpenAIEmbeddings()

if __name__ == '__main__':
    docsearch = Chroma(persist_directory="D:/work/Private/langchain/vector_store", embedding_function=embeddings)
    # 创建问答对象
    qa = RetrievalQA.from_chain_type(llm=OpenAI(), chain_type="stuff", retriever=docsearch.as_retriever(), return_source_documents=True)
    # 进行问答

    def doc_txt(search_query):
        result = qa({"query": search_query})
        return result['result']

    prompts = [
               '江小白有没有衰落',
               ]

    for prompt in prompts:
        print('问:'+prompt+'\n答:'+doc_txt(prompt)+'\n')
